# -*- coding: utf-8 -*-

"""
@Datetime: 2019/1/31
@Author: Zhang Yafei
"""
import os

os.chdir('../')

import sqlite3
from sqlalchemy import create_engine
from Kingadmin.file_reader import reader
from database.MongoDB import MongoPipeline
from application import app


class Database(object):
    """
    将dataframe数据保存到数据库中
    数据库可以是mysql/sqlite3
    """
    def __init__(self):
        for file_name, file_path in app.config.get('FILE_PATH'):
            reader.register(file_name=file_name, file_path=file_path)
        database = list(app.config['DATABASES'].keys())[0]
        if hasattr(self, database):
            self.database = database
            self.conn = getattr(self, database)()

    @staticmethod
    def mysql():
        engine = create_engine(
            app.config.get('DATABASES')['mysql'],
            max_overflow=0,  # 超过连接池大小外最多创建的连接
            pool_size=5,  # 连接池大小
            pool_timeout=30,  # 池中没有线程最多等待的时间，否则报错
            pool_recycle=-1
        )
        return engine

    @staticmethod
    def sqlite3():
        con = sqlite3.connect(app.config.get('DATABASES')['sqlite3'])
        return con

    @staticmethod
    def mongo():
        mongo_uri = app.config.get('DATABASES')['mongo'].get('MONGO_URI')
        db = app.config.get('DATABASES')['mongo'].get('MONGO_DB')
        if mongo_uri:
            return MongoPipeline(mongo_db=db, mongo_uri=mongo_uri)
        return MongoPipeline(mongo_db=db)

    def save(self):
        for file_name,value in reader.enable_admins.items():
            data = value.data
            table_name = file_name.rsplit('.')[0]
            if self.database != 'mongo':
                data.to_sql(name=table_name, con=self.conn, chunksize=1000, if_exists='replace', index=None)
            else:
                self.conn.save(data, collection=table_name)
        if self.database == 'mongo':
            self.conn.close()


if __name__ == '__main__':
    database = Database()
    database.save()